package rabbit.flt.plugins.reactor.http.plugin.impl;

import io.netty.handler.codec.http.HttpResponseStatus;
import org.reactivestreams.Publisher;
import rabbit.flt.common.trace.TraceData;
import rabbit.flt.common.trace.io.HttpResponse;
import rabbit.flt.common.utils.ReflectUtils;
import rabbit.flt.plugins.common.ContextName;
import rabbit.flt.plugins.common.MethodCallback;
import rabbit.flt.plugins.common.WebClientCallBack;
import reactor.core.publisher.Flux;
import reactor.netty.ByteBufFlux;
import reactor.netty.http.client.HttpClient;
import reactor.netty.http.client.HttpClientResponse;

import java.lang.reflect.Method;
import java.util.function.BiFunction;

import static rabbit.flt.plugins.reactor.http.plugin.HttpClientFinalizerPlugin.CONTEXT_KEY_NAME;

public class ResponsePlugin extends ResponseSinglePlugin {

    @Override
    public Object[] before(Object objectEnhanced, Method method, Object[] args) {
        if (isTraceOpened()) {
            args[0] = proxyResponseHandler((BiFunction<HttpClientResponse, ByteBufFlux, Publisher>) args[0]);
        }
        return args;
    }

    /**
     * 代理处理函数
     *
     * @param receiver
     * @return
     */
    private BiFunction<HttpClientResponse, ByteBufFlux, Publisher> proxyResponseHandler(BiFunction<HttpClientResponse, ByteBufFlux, Publisher> receiver) {
        return (httpClientResponse, byteBufFlux) -> Flux.deferContextual(ctx -> {
            TraceData traceData = ctx.get(CONTEXT_KEY_NAME);
            HttpResponse response = traceData.getHttpResponse();
            httpClientResponse.responseHeaders().forEach(entry -> response.addHeader(entry.getKey(), truncate(entry.getValue())));
            response.setStatusCode(httpClientResponse.status().code());
            traceData.updateCost();
            if (HttpResponseStatus.OK.code() != response.getStatusCode()) {
                ByteBufFlux u = ByteBufFlux.fromString(byteBufFlux.aggregate().asString().map(body -> {
                    response.setBody(truncate(body));
                    return body;
                }));
                return receiver.apply(httpClientResponse, u);
            }
            return receiver.apply(httpClientResponse, byteBufFlux);
        });
    }

    /**
     * 调用 response 方法
     *
     * @param objectEnhanced
     * @param method
     * @param callback
     * @param args
     * @return
     */
    @Override
    public Object invokeMethod(Object objectEnhanced, Method method, MethodCallback callback, Object[] args) {
        if (isTraceOpened()) {
            try {
                TraceData traceData = createTraceData(method);
                proxyFlagContext.set(Boolean.TRUE);
                // 发起tcp请求的时候替换原始对象
                HttpClient httpClient = getHttpClient(objectEnhanced, traceData);
                Object result = ReflectUtils.callMethod(httpClient, method, args);
                return Flux.deferContextual(ctx -> {
                    if (!ctx.hasKey(ContextName.WEB_CLIENT_CALL_BACK)) {
                        return ((Flux<Object>) result).contextWrite(context -> context.put(CONTEXT_KEY_NAME, traceData))
                                .doOnComplete(traceData::updateCost)
                                .onErrorResume(handleError(traceData))
                                .doFinally(s -> handleTraceData(s, traceData));
                    } else {
                        // webclient 的next方法导致responseConnection提前cancel
                        WebClientCallBack back = ctx.get(ContextName.WEB_CLIENT_CALL_BACK);
                        back.setCallBack(() -> {
                            traceData.updateCost();
                            handleTraceData(null, traceData);
                        });
                        return ((Flux<Object>) result).contextWrite(context -> context.put(CONTEXT_KEY_NAME, traceData))
                                .onErrorResume(handleError(traceData));
                    }
                });
            } finally {
                proxyFlagContext.remove();
            }
        }
        return super.invokeMethod(objectEnhanced, method, callback, args);
    }

}
